import random

from gxl_ai_utils.utils import utils_file



# remote_file_path,  =  utils_file.do_get_commandline_param(5, ["input_ckpt_path", "output_split_dir", "num_parts", "download_split_ckpt_dir", "download_ckpt_path"])

import os


def do_get_all_available_ckpt_files(remote_root):
    """
    该函数遍历remote_root目录下的所有一级子目录，检查每个子目录中的step_*.pt文件，选择数值最大的文件，并将该文件的路径写入./res.list文件。

    Args:
        remote_root: 需要检查的根目录路径

    Returns:
        None
    """
    total_size = 0

    # 打开文件写入路径
    with open('./res.list', 'w') as res_file:
        # 遍历remote_root下的每个一级子目录
        for subdir in os.listdir(remote_root):
            subdir_path = os.path.join(remote_root, subdir)

            if os.path.isdir(subdir_path):  # 确保是目录
                max_step_file = None
                max_step_value = -1

                # 遍历该子目录下的所有文件
                for file_name in os.listdir(subdir_path):
                    if file_name.startswith('final_') and file_name.endswith('.pt'):
                        max_step_file = os.path.join(subdir_path, file_name)
                        break
                    # 只处理step_*.pt文件
                    if file_name.startswith('step_') and file_name.endswith('.pt'):
                        try:
                            # 提取step_*.pt中的数字部分
                            step_value = int(file_name.split('_')[1].split('.')[0])

                            # 如果该step值大于目前的最大值，则更新
                            if step_value > max_step_value:
                                max_step_value = step_value
                                max_step_file = os.path.join(subdir_path, file_name)
                        except ValueError:
                            # 如果无法转换为整数，跳过
                            continue

                # 如果找到了合法的最大step文件，写入res.list
                if max_step_file:
                    print(max_step_file)
                    res_file.write(max_step_file + '\n')
                    size_i = utils_file.do_get_file_size(max_step_file)
                    total_size += size_i
                    print(max_step_file, size_i/1024, "GB")
    print("total size:", total_size/1024, "GB")


path_list = []


def do_sync_files(file_list_path, password, username, remote_host, local_directory,remote_dir='/'):
    """
    使用 rsync 命令从远程主机下载文件，并支持断点续传(-P决定)。
    remote_dir使用的场景： 如果一个shards_list.txt文件中的pure文件名字不是唯一的， 则需要相对于某一个根目录，保留这个
    目录下面的子目录的结构，这样就不会出现覆盖文件的情况了。
    sync比scp更好用，功能相似的前体下sync支持断点续传
    :param file_list_path: 包含文件列表的文件路径
    :param password: 远程主机的 SSH 密码
    :param username: 远程主机的用户名
    :param remote_host: 远程主机的 IP 地址或主机名
    :param local_directory: 本地存储文件的目标目录
    :param remote_dir: 远程目录起始位置
    :return:
    """
    remote_dir = remote_dir if remote_dir.endswith('/') else remote_dir+'/'
    if len(remote_dir) !=1:
        utils_file.do_replace_str_to_file(remote_dir, '/', file_list_path, file_list_path)
        # file_list_path = './tmp.list'
        # 构造 rsync 命令
        rsync_command = [
            'sshpass', '-p', password,  # 使用 sshpass 提供密码
            'rsync', '-avrP',  # rsync 参数：-a（归档），-v（详细），-r（递归），-P（部分传输和进度）
            # '--no-relative',  # 去掉远程目录的层级结构
            '--ignore-missing-args', # 取消不存在的文件报错
            f'--files-from={file_list_path}',  # 从文件列表中读取文件路径
            f'{username}@{remote_host}:{remote_dir}',  # 远程源路径
            local_directory  # 本地目标目录
        ]
    elif remote_dir=='/':
        # 构造 rsync 命令
        rsync_command = [
            'sshpass', '-p', password,  # 使用 sshpass 提供密码
            'rsync', '-avrP',  # rsync 参数：-a（归档），-v（详细），-r（递归），-P（部分传输和进度）
            '--no-relative',  # 去掉远程目录的层级结构
            '--ignore-missing-args',
            f'--files-from={file_list_path}',  # 从文件列表中读取文件路径
            f'{username}@{remote_host}:{remote_dir}',  # 远程源路径
            local_directory  # 本地目标目录
        ]
    else:
        utils_file.logging_warning('remote dir格式不正确')
        return

    # 执行 rsync 命令
    try:
        utils_file.subprocess.run(rsync_command, check=True)
        print("文件同步完成！")
        utils_file.do_remove_file('./tmp.list')
    except utils_file.subprocess.CalledProcessError as e:
        print(f"rsync 命令执行失败: {e}")
    except FileNotFoundError as e:
        print(f"命令未找到: {e}")
    except Exception as e:
        print(f"发生错误: {e}")


def do_sync_files_download_data_multi_thread(file_list_path, password, username, remote_host,local_directory,remote_dir='/', num_thread=10):
    runner = utils_file.GxlDynamicProcessPool()
    file_list = utils_file.load_list_file_clean(file_list_path)
    file_list_list = utils_file.do_split_list(file_list, num_subsets=num_thread)
    for file_list_i in file_list_list:
        fake_file = utils_file.do_get_fake_file()
        utils_file.write_list_to_file(file_list_i, fake_file)
        runner.add_task(do_sync_files, [fake_file, password, username, remote_host,local_directory,remote_dir])
    runner.start()


if __name__ == '__main__':
    # do_get_all_available_ckpt_files("/mnt/sfs/asr/code/osum_xlgeng_3B/examples/wenetspeech/whisper/exp")
    # do_get_all_available_ckpt_files("/mnt/sfs/asr/ckpt/OSUM_OneEmbedding/examples/osum/exp")
    do_sync_files_download_data_multi_thread(
        file_list_path="./res2.list",
        username="root",
        password="Fy!mATB@QE",
        remote_host="139.210.101.41",
        local_directory="./",
        num_thread=4,
        remote_dir="/mnt/sfs/asr/ckpt/OSUM_OneEmbedding/examples/osum/exp"
    )